from pyvis.network import Network
from neo4j_utils.connect import Neo4jConnection
import os
import random


def visualize_graph():
    # 连接到Neo4j
    neo4j = Neo4jConnection("bolt://127.0.0.1:7687", "neo4j", "abcd1234")

    # 增强的网络配置
    net = Network(
        notebook=False,
        width="100%",
        height="900px",
        bgcolor="#f8f9fa",
        font_color="#333333",
        directed=True,
        cdn_resources="remote"
    )

    # 优化物理引擎参数，避免节点过度拥挤
    net.set_options("""
    var options = {
      "physics": {
        "forceAtlas2Based": {
          "springLength": 350
        },
        "minVelocity": 0.75,
        "solver": "forceAtlas2Based",
        "timestep": 0.35,
        "stabilization": {
          "iterations": 200
        }
      },
      "nodes": {
        "scaling": {
          "min": 15,
          "max": 30,
          "label": {
            "enabled": true
          }
        }
      },
      "edges": {
        "color": {
          "inherit": false
        },
        "smooth": {
          "type": "cubicBezier",
          "forceDirection": "horizontal"
        }
      }
    }
    """)

    try:
        with neo4j.driver.session() as session:
            # 获取所有标签，用于动态颜色配置
            labels_result = session.run("MATCH (n) RETURN DISTINCT labels(n)")
            all_labels = [record[0] for record in labels_result if record[0]]

            # 为每个标签生成独特颜色
            label_colors = {}
            for labels in all_labels:
                for label in labels:
                    if label not in label_colors:
                        # 生成柔和的随机颜色
                        r = random.randint(50, 200)
                        g = random.randint(50, 200)
                        b = random.randint(50, 200)
                        label_colors[label] = f"rgb({r}, {g}, {b})"

            # 读取节点并添加到网络
            nodes = session.run("MATCH (n) RETURN elementId(n) as id, labels(n) as labels, properties(n) as props")
            for node in nodes:
                node_id = node["id"]
                labels = node["labels"]
                props = node["props"]

                # 确定节点颜色
                if labels:
                    color = label_colors.get(labels[0], "#3498db")
                else:
                    color = "#95a5a6"  # 默认灰色

                # 节点标题（悬停时显示）
                title = f"<b>{props.get('name', f'Node {node_id}')}</b><br>"
                title += "<br>".join([f"{k}: {v}" for k, v in props.items()])

                # 根据标签设置节点大小和形状
                size = 25 if "Country" in labels else 20  # 国家节点稍大
                shape = "ellipse" if "Country" in labels else "circle"

                net.add_node(
                    node_id,
                    label=props.get("name", f"Node {node_id}"),
                    color=color,
                    title=title,
                    size=size,
                    shape=shape
                )

            # 读取关系并添加到网络
            relationships_types = session.run("MATCH ()-[r]->() RETURN DISTINCT type(r)")
            rel_types = [record[0] for record in relationships_types]

            # 关系类型中英文映射
            rel_type_mapping = {
                "HAS_CAPITAL": "首都"
                # 可以添加其他关系的映射
            }

            # 为不同关系类型设置不同颜色
            rel_colors = {}
            for rel_type in rel_types:
                if rel_type not in rel_colors:
                    r = random.randint(100, 200)
                    g = random.randint(100, 200)
                    b = random.randint(100, 200)
                    rel_colors[rel_type] = f"rgba({r}, {g}, {b}, 0.7)"

            # 只获取每种关系一次，避免重复
            relationships = session.run(
                "MATCH (a)-[r]->(b) RETURN DISTINCT elementId(a) as source, elementId(b) as target, type(r) as type, properties(r) as props")
            for rel in relationships:
                # 使用中文标签显示
                display_label = rel_type_mapping.get(rel["type"], rel["type"])

                # 关系标题
                rel_title = f"<b>{display_label}</b>"
                if rel['props']:
                    rel_title += "<br>" + "<br>".join([f"{k}: {v}" for k, v in rel['props'].items()])

                net.add_edge(
                    rel["source"],
                    rel["target"],
                    label=display_label,  # 使用映射后的标签
                    color=rel_colors.get(rel["type"], "#999999"),
                    title=rel_title,
                    width=2,
                    arrows="to"
                )

        # 保存可视化结果
        output_file = "enhanced_knowledge_graph.html"
        net.write_html(output_file)

        # 验证文件是否生成
        if os.path.exists(output_file):
            print(f"优化后的知识图谱已保存至: {os.path.abspath(output_file)}")
        else:
            raise Exception(f"未能生成文件: {output_file}")

    except Exception as e:
        print(f"可视化过程出错: {str(e)}")
    finally:
        neo4j.close()
